from torch.utils.data import DataLoader
from torchvision.datasets import FashionMNIST
from torchvision import transforms
from torchvision.datasets import ImageFolder
from torch.utils import data as Data
"""
FashionMNIST 数据集是28x28 大小，单通道
AlexNet 网络的输入是227x227的大小
"""


# 获取数据加载方法
class GetDataLoader:
    # 图片大小
    IMG_SIZE = (227,227)
    # 图片的通道
    IMG_CHANNELS = 1
    # 单例模式：对象的地址信息
    INSTANCE = None
    def __new__(cls, *args, **kwargs):
        if cls.INSTANCE:
            return cls.INSTANCE
        else:

            cls.INSTANCE = super().__new__(cls)
            return cls.INSTANCE

    def __init__(self):
        # 数据集的跟目录
        self.data_root = '../a_FashionMNIST'
        # 图片的种类
        self.IMG_CLASSES:list = []
        # 训练时：训练集和验证集
        self.train_and_val_dataloader:tuple = self._make_train_and_val_dataloader()
        # 验证时：测试数据集
        self.test_dataloader = self._make_test_dataloader()

    # 训练时：训练数据集和验证数据集
    def _make_train_and_val_dataloader(self):
        train_data = FashionMNIST(root=self.data_root,
                                  train=True,
                                  transform=transforms.Compose(
                                      [transforms.Resize(self.IMG_SIZE), transforms.ToTensor()]),
                                  download=True
                                  )
        # 设置好图片的种类
        self.IMG_CLASSES = train_data.classes
        # 数据集分割
        train_data, val_data = Data.random_split(
            train_data,
            [round(0.8 * len(train_data)), round(0.2 * len(train_data))]
        )

        # 训练数据加载
        train_dataloader = Data.DataLoader(
            dataset=train_data,
            batch_size=64,
            shuffle=True,
        )

        # 验证数据集加载
        val_dataloader = Data.DataLoader(
            dataset=val_data,
            batch_size=64,
            shuffle=True
        )

        return train_dataloader, val_dataloader

    # 测试时：测试使用的数据集
    def _make_test_dataloader(self):
        test_data = FashionMNIST(
            root=self.data_root,
            train=False,
            transform=transforms.Compose([transforms.Resize(self.IMG_SIZE), transforms.ToTensor()]),
            download=True,
        )

        test_dataloader = DataLoader(
            dataset=test_data, batch_size=1, shuffle=True, num_workers=0
        )

        return test_dataloader
if __name__ == '__main__':
    # 数据集的情况
    loader = GetDataLoader()
    # 数据集中图片的通道
    channels = loader.IMG_CHANNELS
    # 图片要resize的大小
    img_size = loader.IMG_SIZE
    # 图片的类型，[分类名1,分类名2]
    img_classes = loader.IMG_CLASSES
    # 训练数据加载和验证数据加载
    train_dataloader, val_dataloader = loader.train_and_val_dataloader
    # 测试数据加载
    test_dataloader = loader.test_dataloader